package com.llseven.ai.spring.llseven.service;

import com.llseven.ai.spring.llseven.enumtype.ModelType;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.web.client.RestTemplateBuilder;
import org.springframework.http.*;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Slf4j
@Service
public class BailianChatService {
    
    private final RestTemplate restTemplate;
    private final HttpHeaders defaultHeaders;
    private final String apiKey = "sk-89826b76d3384f3e9d6329186e1b3471";
    
    public BailianChatService(RestTemplateBuilder builder) {
        this.restTemplate = builder.build();
        this.defaultHeaders = new HttpHeaders();
        this.defaultHeaders.setContentType(MediaType.APPLICATION_JSON);
        this.defaultHeaders.set("Authorization", "Bearer " + apiKey);
    }
    
    public String chat(String prompt) {
        String url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation";
        
        Map<String, Object> body = new HashMap<>();
        body.put("model", "qwen-turbo");
        
        Map<String, Object> input = new HashMap<>();
        input.put("prompt", prompt);
        body.put("input", input);
        
        Map<String, Object> params = new HashMap<>();
        params.put("temperature", 0.8);
        body.put("parameters", params);
        
        HttpEntity<Map<String, Object>> entity = new HttpEntity<>(body, defaultHeaders);
        
        ResponseEntity<Map> response = restTemplate.postForEntity(url, entity, Map.class);
        Map<String, Object> output = (Map<String, Object>) response.getBody().get("output");
        return output != null ? output.get("text").toString() : "No response";
    }
    
    
    public String create(ModelType modelType, String prompt, String negativePrompt, String refImgUrl, Double strength) {
        String url;
        Map<String, Object> body = new HashMap<>();
        Map<String, Object> input = new HashMap<>();
        Map<String, Object> parameters = new HashMap<>();
        
        switch (modelType) {
            case WANX_T2I_TURBO:
            case SD_3_5_LARGE:
                url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis";
                body.put("model", modelType.getValue());
                input.put("prompt", prompt);
                input.put("init_image", refImgUrl);
                parameters.put("n", 1);
                parameters.put("seed", 35);
                parameters.put("cfg", strength);
                break;
            case WANX_V1:
                url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis";
                body.put("model", modelType.getValue());
                input.put("prompt", prompt);
                input.put("negative_prompt", negativePrompt);
                input.put("ref_img", refImgUrl);
                parameters.put("n", 1);
                parameters.put("prompt_extend", false);
                parameters.put("ref_strength", strength);
                break;
            
            case WANX_T2I_IMAGEEDIT:
                url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2image/image-synthesis";
                body.put("model", modelType.getValue());
                input.put("prompt", prompt);
                input.put("base_image_url", refImgUrl);
                input.put("function", "description_edit");
                parameters.put("n", 1);
                parameters.put("strength", strength);
                parameters.put("watermark", true);
                break;
//            case "wanx-x-painting":
//                url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2image/image-synthesis";
//                body.put("model", modelType);
//                input.put("prompt", prompt);
//                input.put("base_image_url", refImgUrl);
//                input.put("mask_image_url", "http://suwchu9pz.hn-bkt.clouddn.com/ai/mask/202504191124142.png");
//                input.put("function", "description_edit_with_mask");
//                parameters.put("n", 1);
//                break;
            
            default:
                return "Unsupported model type: " + modelType;
        }
        
        body.put("input", input);
        body.put("parameters", parameters);
        
        HttpHeaders headers = new HttpHeaders(defaultHeaders);
        headers.set("X-DashScope-Async", "enable");
        
        HttpEntity<Map<String, Object>> entity = new HttpEntity<>(body, headers);
        ResponseEntity<Map> response = restTemplate.postForEntity(url, entity, Map.class);
        Map<String, Object> output = (Map<String, Object>) response.getBody().get("output");
        return output != null ? output.get("task_id").toString() : "No task ID returned";
    }
    
    public String query(String taskId) {
        String url = "https://dashscope.aliyuncs.com/api/v1/tasks/" + taskId;
        
        HttpHeaders headers = new HttpHeaders(defaultHeaders);
        HttpEntity<String> entity = new HttpEntity<>(headers);
        
        try {
            ResponseEntity<Map> response = restTemplate.exchange(url, HttpMethod.GET, entity, Map.class);
            
            if (response.getStatusCode().is2xxSuccessful()) {
                Map<String, Object> responseBody = response.getBody();
                if (responseBody != null) {
                    Map<String, Object> output = (Map<String, Object>) responseBody.get("output");
                    if (output != null && "SUCCEEDED".equals(output.get("task_status"))) {
                        List<Map<String, Object>> results = (List<Map<String, Object>>) output.get("results");
                        if (results != null && !results.isEmpty()) {
                            return results.get(0).get("url").toString();
                        }
                    }
                }
            }
        } catch (Exception e) {
            log.error("Query failed for taskId {}: {}", taskId, e.getMessage());
            return "Error occurred: " + e.getMessage();
        }
        
        return "Task not completed or failed";
    }
}
